#ifndef PARTITIONING_HPP
#define PARTITIONING_HPP

#include "typeAliases.hpp"
#include "meshTypes_fwd.hpp"
#include "mpiHelpers.hpp"
#include "structMeshPart.hpp"
#include "structMeshOverlap.hpp"
#include <array>
#include <iostream>
#include <tuple>
#include <vector>

namespace hms {

/**
 * @brief Performs integer factorisation.
 * 
 * @param num The integer to be factorised.
 * @return std::vector<int> The arguments' factors.
 * 
 * Integer factorisation is a difficult problem. The algorithm implemented here
 * is intended for small integers, as it is a somewhat naive implementation.
 * Returns a vector with length 1, if the argument is the number 1, and a 
 * vector with length 2 if it is a prime number. Triggers assertion if zero or
 * negative number are passed in.
 */
std::vector<int> factors(int num);

/**
 * @brief "Rounds" a number to the nearest match in the provided vector.
 * Intended for factors, but will work with any vector that is sorted in
 * ascending order.
 * 
 * @param num The number to be rounded.
 * @param factors The numbers among which the closest match is to be found.
 * Must be sorted in ascending order.
 * @return int The closest match of num in factors.
 */
int roundToFactor(double num, const std::vector<int>& factors);

/**
 * @brief Takes the number of cells in topological x- and y-direction as well
 * as the number of nodes/ranks and returns the optimal number of partitions
 * in x- and y-direction.
 * 
 * @param nCellsX The number of cells in (topological) x-direction.
 * @param nCellsY The number of cells in (topological) y-direction.
 * @param nodes The number of nodes/ranks.
 * @return Array2i The number of nodes in x- and y-direction.
 * 
 * "Optimal" here means that the resulting mesh partitions are as compact as
 * possible and thus have minimal interfaces.
 */
Array2i structMeshNParts(Index nCellsX, Index nCellsY, int nodes);

/**
 * @overload
 */
Array2i structMeshNParts( const Ref<const Array2i>& meshSizeGlobal, int nodes );

Array2i structMeshPartSize(
	const Ref<const Array2i>& meshSizeGlobal,
	const Ref<const Array2i>& nParts,
	const MPI_Comm_H& mpi
);

/**
 * @brief Calculates the offset of the caller's rank's mesh partition, based
 * on the assumption of equal sizes, i.e. the initial offsets before any
 * balancing.
 * 
 * @param meshSizeGlobal 
 * @param nParts 
 * @param comm 
 * @return Array2i
 * 
 * If the mesh sizes are not equal, use meshPartSizesToOffsets() instead.
 */
Array2i structMeshInitOffset(
	const Ref<const Array2i>& meshSizeGlobal,
	const Ref<const Array2i>& nParts,
	const MPI_Comm_H& comm
);

/**
 * @brief Takes an array and calculates the length of n segments with similar
 * partial sums.
 * 
 * @param loads The input array.
 * @param nSegments The number of segments.
 * @return ArrayXi An array of length n, which contains the segment lengths.
 */
ArrayXi equalPartialSums(
	const Ref<const ArrayXs>& loads,
	Index nSegments
);


/**
 * @brief Takes the number of partitions in x- and y-direction and all 
 * mesh partition sizes and returns the offsets from the bottom left corner
 * domain cell for all mesh partitions.
 * 
 * @param nParts The number of partitions in x- and y-dimension.
 * @param meshPartSizes All mesh partition sizes.
 * @return Array2Xi All mesh partition offsets
 */
Array2Xi meshPartSizesToOffsets(
	const Ref<const Array2i>& nParts,
	const Ref<const Array2Xi>& meshPartSizes
);


/**
 * @brief Calculates mesh sizes with similar loads. The load data argument
 * must be of the same length as the global mesh size.
 * 
 * @param meshSizeGlobal The mesh size in x- and y-dimension.
 * @param nParts The number of partitions in x- and y-dimension.
 * @param loadPerCell The per-cell load data.
 * @return Array2Xi Balanced partition sizes (x- and y-dimension).
 * 
 * Use this version if you have per-cell load data collected on one node, where
 * the new part sizes are calculated locally.
 */
Array2Xi balancedStructMeshParts(
	const Ref<const Array2i>& meshSizeGlobal,
	const Ref<const Array2i>& nParts,
	const Ref<const ArrayXs>& loadPerCell
);

/**
 * @brief Calculates mesh sizes with similar loads and returns them to the
 * part's comm's master (i.e. does not broadcast sizes).
 * 
 * @param part 
 * @param localLoadData The local per-cell load data as a flattened array.
 * Its number of elements must be equal to the number of domain cells of that
 * node's mesh partition.
 * @return Array2Xi Mesh sizes for all partitions. Each column corresponds to
 * one node, with the first row containing the size in x- and the second row
 * the size in y-direction.
 * 
 * Use this version when you have local per-cell load data on each node.
 */
template<typename MeshType>
Array2Xi balancedStructMeshParts( 
	const StructMeshPart<MeshType>& part,
	const ArrayXs& localLoadData
);

/**
 * @brief Calculates mesh sizes with similar loads. This version is for a single
 * load value per node, e.g. the CPU load as reported by the OS.
 * 
 * @param meshSizeGlobal The mesh size in x- and y-dimension.
 * @param nParts The number of partitions in x- and y-dimension.
 * @param meshSizes The size of all mesh partitions.
 * @param loads The per-node load data.
 * @return Array2Xi Balanced partition sizes (x- and y-dimension).
 * 
 * Use this version if you have per-node load data collected on one node.
 */
Array2Xi balancedStructMeshParts(
	const Ref<const Array2i>& meshSizeGlobal,
	const Ref<const Array2i>& nParts,
	const Ref<const Array2Xi>& meshSizes,
	const Ref<const ArrayXs>& loadPerNode
);



/**
 * @brief Distributes an array from the specified root to all nodes.
 * 
 * @tparam Derived1 
 * @tparam Derived2 
 * @param globalArray The global array from which the values are read and
 * distributed. On root, this must have the correct size.
 * @param localArray The local array from which values are read and written to
 * the global array.
 * @param nComponents The number of values per cell.
 * @param globalArraySize Basically the global mesh size.
 * @param localArraySize Basically the local mesh size.
 * @param localOffset The local mesh offset.
 * @param comm 
 * @param root 
 * 
 * Use this version if your local array already exists.
 */
template<typename Derived1, typename Derived2>
void decomposeArray(
	const Eigen::DenseBase<Derived1>& globalArray,
	Eigen::DenseBase<Derived2>& localArray,
	Index nComponents,
	const Ref<const Array2i>& globalArraySize,
	const Ref<const Array2i>& localArraySize,
	const Ref<const Array2i>& localOffset,
	const MPI_Comm_H& comm = MPI_COMM_WORLD,
	int root = 0
);


/**
 * @brief For internal use.
 */
template<typename Derived>
auto decomposeArray(
	const Eigen::DenseBase<Derived>& array,
	const Ref<const Array2i>& globalArraySize,
	const Ref<const Array2i>& localArraySize,
	const Ref<const Array2i>& localOffset,
	const MPI_Comm_H& comm = MPI_COMM_WORLD,
	int root = 0
) -> typename Derived::PlainObject;

/**
 * @brief Distributes an array from the specified root to all nodes.
 * 
 * @tparam Derived A type derived from Eigen::DenseBase
 * @param array The array to be distributed. Data must be provided on root.
 * @param part StructMeshPart which holds relevant data.
 * @param withGhosts Switch, which allows to distribute arrays that contain
 * ghost cell data.
 * @param root The rank which holds the yet undistributed array.
 * @return Derived Same type as the passed-in array.
 * 
 * The size of the array on root must conform to the global mesh size plus
 * ghost layers, if withGhosts is true.
 * If a single value per cell is used, the array must be a row or column vector.
 * If multiple values per cell are used, the array must have as many rows as
 * values per cell and as many columns as the global mesh size (plus ghost
 * cells, if so specified).
 */
template<typename Derived, typename MeshType>
auto decomposeArray(
	const Eigen::DenseBase<Derived>& array,
	const StructMeshPart<MeshType>& part,
	bool withGhosts = false,
	int root = 0
) -> typename Derived::PlainObject;


/**
 * @brief Recomposes an array from all nodes on the specified root.
 * @overload
 * 
 * @tparam Derived1 
 * @tparam Derived2 
 * @param globalArray The global array into which the values are to be
 * recomposed. On root, this must have the correct size.
 * @param localArray The local array from which values are read and written to
 * the global array.
 * @param nComponents The number of values per cell.
 * @param globalArraySize Basically the global mesh size.
 * @param localArraySize Basically the local mesh size.
 * @param localOffset The local mesh offset.
 * @param localGhostOffBeg The local ghost cells to be skipped at the beginning.
 * @param localGhostOffEnd The local ghost cells to be skipped at the end.
 * @param comm Communicator
 * @param root 
 * @return void 
 * 
 * This overload takes the global array as a parameter and requires it to have
 * the correct size. Use this if your global array already exists.
 */
template<typename Derived1, typename Derived2>
void recomposeArray(
	Eigen::DenseBase<Derived1>& globalArray,
	const Eigen::DenseBase<Derived2>& localArray,
	Index nComponents,
	const Ref<const Array2i>& globalArraySize,
	const Ref<const Array2i>& localArraySize,
	const Ref<const Array2i>& localOffset,
	const Ref<const Array2i>& localGhostOffBeg = Array2i{0,0},
	const Ref<const Array2i>& localGhostOffEnd = Array2i{0,0},
	const MPI_Comm_H& comm = MPI_COMM_WORLD,
	int root = 0
);


/**
 * @brief Recomposes an array from all nodes on the specified root.
 * @overload
 * 
 * If possible, use the overload taking a StructMeshPart.
 */
template<typename Derived>
auto recomposeArray(
	const Eigen::DenseBase<Derived>& localArray,
	const Ref<const Array2i>& globalArraySize,
	const Ref<const Array2i>& localArraySize,
	const Ref<const Array2i>& localOffset,
	const Ref<const Array2i>& localGhostOffBeg = Array2i{0,0},
	const Ref<const Array2i>& localGhostOffEnd = Array2i{0,0},
	const MPI_Comm_H& comm = MPI_COMM_WORLD,
	int root = 0
) -> typename Derived::PlainObject;

/**
 * @brief Recomposes an array from all nodes on the specified root.
 * 
 * @tparam Derived A type derived from Eigen::DenseBase.
 * @param localArray The local arrays to be recomposed.
 * @param part StructMeshPart which holds relevant data.
 * @param withGhosts Switch, which allows to recompose arrays that contain
 * ghost cell data.
 * @param useOldSizes Switch, which allows the local arrays to conform to the 
 * old mesh size, before the last rebalancing.
 * @param root Rank of the node on which the local arrays are recomposed.
 * @return Derived Same type as the passed-in local arrays.
 * 
 * @see decomposeArray
 */
template<typename Derived, typename MeshType>
auto recomposeArray(
	const Eigen::DenseBase<Derived>& localArray,
	const StructMeshPart<MeshType>& part,
	bool withGhosts = false,
	bool useOldSizes = false,
	int root = 0
) -> typename Derived::PlainObject;


/**
 * @brief Uses offsets in one dimension to find the partition in which a cell
 * lies.
 * 
 * @param cell Cell index of the same dimension as the offsets
 * @param offsets 
 * @return int Partition index
 */
int findPart( Index cell, const Ref<const ArrayXi>& offsets );

/**
 * @brief For two partitions, provided as their index pairs consisting of
 * band and index within band, returns the overlap between the two as minimum
 * and maximum cell indices in x and y.
 * 
 * @param part Index pair for current partition as {Index in band, band}
 * @param targetPart Index pair for target partition as {Index in band, band}
 * @param nParts Number of partitions in x- and y-dimensions
 * @param ownOffset 
 * @param ownSize 
 * @param targetOffset 
 * @param targetSize 
 * @param ng Number of ghost cell layers
 * @return std::optional<ArrayNNi<2,2>> Returns std::nullopt if no overlap
 * exists. Otherwise, returns the overlapping block as:
 * xMin  xMax
 * yMin  yMax
 * 
 * @see overlappingBlocks()
 */
std::optional<ArrayNNi<2,2>> overlappingBlock(
	const Ref<const Array2i>& part,
	const Ref<const Array2i>& targetPart,
	const Ref<const Array2i>& nParts,
	const Ref<const Array2i>& ownOffset,
	const Ref<const Array2i>& ownSize,
	const Ref<const Array2i>& targetOffset,
	const Ref<const Array2i>& targetSize,
	int ng
);

/**
 * @brief Returns the overlapping blocks between a partition's field-like data
 * and all other partition's respective data.
 * Intended to be used with ownOffset and ownSize from after rebalancing and 
 * targetOffsets and targetSizes from before rebalancing.
 * 
 * @param part Index pair representing the position of the executing partition
 * in the global mesh. Consists of {Index in band, band}.
 * @param nParts Number of partitions in x- and y-direction. nParts.y() is the
 * number of bands.
 * @param ownOffset Index pair of own bottom left cell within global mesh.
 * @param ownSize Number of cells in x- and y-direction.
 * @param targetOffsets All partitions' offsets.
 * @param targetSizes All partitions' sizes.
 * @param ng The number of ghost cell layers that the field-like data contains.
 * @return std::vector<OverlapBlock> 
 * 
 * @see overlappingBlock()
 */
std::vector<OverlapBlock> overlappingBlocks(
	const Ref<const Array2i>& part,
	const Ref<const Array2i>& nParts,
	const Ref<const Array2i>& ownOffset,
	const Ref<const Array2i>& ownSize,
	const Ref<const Array2Xi>& targetOffsets,
	const Ref<const Array2Xi>& targetSizes,
	int ng
);


/**
 * @brief Redistributes values from old to new mesh sizes.
 * 
 * @param newValues The array expression into which the redistributed values are
 * to be written.
 * @param values The values to be redistributed. Number of rows must equal the
 * number of values per cell and the number of columns must equal the number of
 * cells in the old mesh size, including potential ghost cells.
 * E.g. Field::values() works.
 * @param meshSize The new mesh size.
 * @param offset The new mesh offset, i.e. the x- and y-indices of its bottom
 * left domain cell within the global mesh.
 * @param oldSizes All old mesh sizes.
 * @param oldOffsets All old offsets.
 * @param overlappingBlocks The return value from hms::overlappingBlocks().
 * Taken as parameter to allow reusage.
 * @param comm The communicator.
 * 
 * This version takes an array expression and writes the redistributed values
 * to it. This allows block expressions to be used.
 */
void redistributeValues(
	Ref<ArrayXXs> newValues,
	const Ref<const ArrayXXs>& values,
	const Ref<const Array2i>& meshSize,
	const Ref<const Array2i>& offset,
	const Ref<const Array2Xi>& oldSizes,
	const Ref<const Array2Xi>& oldOffsets,
	const std::vector<OverlapBlock>& overlappingBlocks,
	const MPI_Comm_H& comm
);

/**
 * @brief Redistributes values from old to new mesh sizes.
 * @overload
 * 
 * @return ArrayXXs The values pertaining to the new size.
 * 
 * This version internally creates an array of the right size to hold the
 * redistributed values, fills and returns it.
 */
ArrayXXs redistributeValues(
	const Ref<const ArrayXXs>& values,
	const Ref<const Array2i>& meshSize,
	const Ref<const Array2i>& offset
	const Ref<const Array2Xi>& oldSizes,
	const Ref<const Array2Xi>& oldOffsets,
	const std::vector<OverlapBlock>& overlappingBlocks,
	const MPI_Comm_H& comm
);


}

#include "structMeshPartitioning.tpp"

#endif